#include "co.h"
#include <stdlib.h>
#include <setjmp.h>
#include <stdint.h>
#include <assert.h>
#include <time.h>

// #define DEBUG
#define STACK_SIZE 64*1024
#define CO_NUM 128
#define false 0
#define true 1

#include <stdio.h>

#define Assert(expr, str) if((expr)==false) {printf(str);assert(expr);} else ;

/* ========== Data Types and Structs ========== */
enum co_status{
  CO_NEW = 1, 
  CO_RUNNING, 
  CO_WAITING, 
  CO_DEAD, 
};

struct co {
  const char* name;
  void (*func)(void*);
  void* arg;

  enum co_status status;
  struct co* waiter;
  jmp_buf context;
  uint8_t stack[STACK_SIZE];
  struct co* next;
};

typedef struct co co;
typedef void(*pf)(void (*func)(void*), void* arg, co* this);

/* ========== externs ========== */


/* ========== Static ========== */
int _co_count = 0;
struct co* _current;
struct co* HEAD;
struct co* TAIL;

/* ========== Functions ========== */
void wrapper(void (*func)(void*), void* arg, co* this) {
  this->status = CO_RUNNING;
  ((void(*)())func) (arg);
#ifdef DEBUG
  printf("-----------------------------\nfinish %s\n------------------------------\n", this->name);
#endif
  this->status = CO_DEAD;
  if (this->waiter) this->waiter->status = CO_RUNNING;
  co_yield();
}

void __attribute__((constructor)) start() {
  _co_count = 1;
  _current = malloc(sizeof(co));
  Assert(_current, "init co failed!\n");
  HEAD = TAIL = _current;
  _current->next = HEAD;
  _current->name = "main";
  _current->status = CO_RUNNING;
  srand((int)time(0));
}

struct co *co_start(const char *name, void (*func)(void *), void *arg) {
  co * new = malloc(sizeof(co));
  Assert(new, "malloc co failed!\n");
  new->name = name;
  new->func = func;
  new->arg = arg;
  new->status = CO_NEW;
  new->waiter = NULL;
  new->next = HEAD;
  HEAD = new;
  TAIL->next = new;
  _co_count += 1;
#ifdef DEBUG
  printf("Create co, name %s, co count is %d\n", new->name, _co_count);
#endif
  return new;
}

void co_wait(struct co *co) {
  if (co->status==CO_DEAD) {
    if (HEAD==co) HEAD = co->next;
    else {
      struct co* tmp = HEAD;
      while(tmp->next!=co) {
        tmp = tmp->next;
      }
      tmp->next = co->next;
      if (TAIL==co) TAIL = tmp;
    }

    free(co);
    _co_count--;
    return;
  }
  co->waiter = _current;
  _current->status = CO_WAITING;
  co_yield();

  Assert((_current->status)!=(CO_WAITING), "_current is waiting but yield to!\n");

  if (HEAD==co) HEAD = co->next;
  else {
    struct co* tmp = HEAD;
    while(tmp->next!=co) {
      tmp = tmp->next;
    }
    tmp->next = co->next;
    if (TAIL==co) TAIL = tmp;
  }
  free(co);
  _co_count--;
  return;
  
}

void co_yield() {
  assert(_current);
  // save buf
  int ret;
  if ((ret = setjmp(_current->context)) != 0) return;

  // ret == 0, save buf and yield 
  int ind = rand()%(_co_count);
  // printf("------%d\n", ind);
  co* c = HEAD;
  while (ind>0 || (c->status!=CO_RUNNING && c->status!=CO_NEW)) {
    c = c->next;
    ind--;
    // if (ind>-100) {
    //   printf("%s ", c->name);
    //   printf("%d ", ind);
    //   printf("%d\n", c->status);
    // }
  }
  Assert(c, "No co is chosen to yield to!\n");

#ifdef DEBUG
  printf("yield to %s, status is %d\n", c->name, c->status);
#endif
  _current = c;
  if (c->status==CO_NEW) {
    // printf("before start!\n");
    // asm volatile (
    // #if __x86_64__
    //   "movq %0, %%rsp; movq %2, %%rdi; jmp *%1" : : "b"((uintptr_t)c->stack+sizeof(c->stack)),   "d"(c->func), "a"(c->arg)
    // #else
    //   "movl %0, %%esp; movl %2, 4(%0); jmp *%1" : : "b"((uintptr_t)c->stack-8+sizeof(c->stack)), "d"(c->func), "a"(c->arg)
    // #endif
    // );
  #if __x86_64__
    asm (
        "movq %[sp], %%rsp\n"
        // "movq %[func], %%rdi\n"
        // "movq %[arg], %%rsi\n"
        // "movq %[cp], %%rdx\n"
        // "jmp *%[wrapper]\n"
        :
        :[sp]"r"((uintptr_t)c->stack+sizeof(c->stack))/*, [func]"r"(c->func), [arg]"r"(c->arg), [cp]"r"(c), [wrapper]"r"(wrapper)*/
        :"rdi", "rsi", "rdx"
    );
    wrapper(c->func, c->arg, c);
  #else
    asm (
        "movl %[sp], %%esp\n"
        "movl %[func], 4(%[sp])\n"
        "movl %[arg], 8(%[sp])\n"
        "movl %[cp], 0xc(%[sp])\n"
        "jmp *%[wrapper]\n"
        :
        :[sp]"r"((uintptr_t)c->stack+sizeof(c->stack)-0x10), [func]"r"(c->func), [arg]"r"(c->arg), [cp]"r"(c), [wrapper]"r"(wrapper)
        :
    );
  #endif


  } else if (c->status==CO_RUNNING) {
  #ifdef DEBUG
    printf("  retoring %s\n", c->name);
  #endif
    longjmp(c->context, 1);
  } else {
    Assert(0, "c->status error!\n");
  }
}
